import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn import Parameter

from .misc import ModuleWrapper


def calculate_kl(log_alpha):
    return 0.5 * torch.sum(torch.log1p(torch.exp(-log_alpha)))

class BBBConv2d(ModuleWrapper):
    
    def __init__(self, in_channels, out_channels, kernel_size, stride=1,
                 padding=0, dilation=1, bias=True, groups=1, alpha_shape=(1,1), name='BBBConv2d'):
        super(BBBConv2d, self).__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = (kernel_size, kernel_size)
        self.stride = stride
        self.padding = padding
        self.dilation = dilation
        self.alpha_shape = alpha_shape
        self.groups = groups
        self.weight = Parameter(torch.Tensor(
            out_channels, in_channels // groups, *self.kernel_size))
        if bias:
            self.bias = Parameter(torch.Tensor(1, out_channels, 1, 1))
        else:
            self.register_parameter('bias', None)
        self.out_bias = lambda input, kernel: F.conv2d(input, kernel, self.bias, self.stride, self.padding, self.dilation, self.groups)
        self.out_nobias = lambda input, kernel: F.conv2d(input, kernel, None, self.stride, self.padding, self.dilation, self.groups)
        self.log_alpha = Parameter(torch.Tensor(*alpha_shape))
        self.reset_parameters()
        self.name = name

    def reset_parameters(self):
        n = self.in_channels
        for k in self.kernel_size:
            n *= k
        stdv = 1. / math.sqrt(n)
        self.weight.data.uniform_(-stdv, stdv)
        if self.bias is not None:
            self.bias.data.uniform_(-stdv, stdv)
        self.log_alpha.data.fill_(-5.0)

    def forward(self, x):

        mean = self.out_bias(x, self.weight)

        if self.training:
            sigma = torch.exp(self.log_alpha) * self.weight.pow(2)
            std = torch.sqrt(1e-16 + self.out_nobias(x.pow(2), sigma))
            epsilon = std.data.new(std.size()).normal_()
            out = mean + std * epsilon
        else:
            out = mean

        return out

    def kl_loss(self):
        return [calculate_kl(self.log_alpha)]



